use rand::Rng;

#[inline(always)]
pub fn is_prime(n: u128) -> bool {
    const NUM_BASES: usize = 14;
    #[rustfmt::skip]
    const WITNESSES: [(u128, &[u128]); NUM_BASES] = [
        (2_046,                             &[2]                                                        ),
        (1_373_652,                         &[2, 3]                                                     ),
        (9_080_190,                         &[31, 73]                                                   ),
        (25_326_000,                        &[2, 3, 5]                                                  ),
        (4_759_123_140,                     &[2, 7, 61]                                                 ),
        (1_112_004_669_632,                 &[2, 13, 23, 1_662_803]                                     ),
        (2_152_302_898_746,                 &[2, 3, 5, 7, 11]                                           ),
        (3_474_749_660_382,                 &[2, 3, 5, 7, 11, 13]                                       ),
        (341_550_071_728_320,               &[2, 3, 5, 7, 11, 13, 17]                                   ),
        (3_825_123_056_546_413_050,         &[2, 3, 5, 7, 11, 13, 17, 19, 23]                           ),
        (u64::MAX as u128,                  &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]               ),
        (318_665_857_834_031_151_167_461,   &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]               ),
        (3_317_044_064_679_887_385_961_981, &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]           ),
        (u128::MAX,                         &[2, 325, 9375, 28_178, 450_775, 9_780_504, 1_795_265_022]  ),
    ];

    if n == 2 || n == 3 {
        return true;
    } else if n <= 1 || n % 2 == 0 || n % 3 == 0 {
        return false;
    }

    let mut i = 5;
    let l = n.ilog2() as u128;
    while i <= l {
        if n % i == 0 || n % (i + 2) == 0 {
            return false;
        }
        i += 6;
    }

    let mut d = n - 1;
    while d % 2 == 0 {
        d >>= 1;
    }

    let mut i = 0;
    while i < NUM_BASES && WITNESSES[i].0 < n {
        i += 1;
    }
    let witnesses = WITNESSES[i].1;

    let mut i = 0;
    while i < witnesses.len() && witnesses[i] < n {
        if !miller_test(d, n, witnesses[i]) {
            return false;
        }
        i += 1;
    }

    true
}

#[inline(always)]
fn miller_test(mut d: u128, n: u128, k: u128) -> bool {
    let mut x = pow_mod(k, d, n);
    if x == 1 || x == n - 1 {
        return true;
    }

    while d != n - 1 {
        x = mul_mod(x, x, n);
        d *= 2;

        if x == 1 {
            return false;
        } else if x == n - 1 {
            return true;
        }
    }

    false
}

#[inline(always)]
pub fn pollard_rho(p: u128) -> u128 {
    let mut rng = rand::thread_rng();
    loop {
        let mut x = rng.gen_range(0..p);
        let mut y = x;
        let mut z = 1;
        let c = rng.gen_range(0..p);
        let (mut i, mut j) = (0, 1);
        loop {
            x = add_mod(mul_mod(x, x, p), c, p);
            z = mul_mod(z, y.abs_diff(x), p);
            if x == y || z == 0 { break }
            if i % 127 == 0 || i == j {
                let g = gcd(z, p);
                if g > 1 { return g }
                if i == j {
                    y = x;
                    j <<= 1;
                }
            }
            i += 1;
        }
    }
}

#[inline(always)]
pub fn pow_mod(mut base: u128, mut exp: u128, modulo: u128) -> u128 {
    let mut res = 1;

    base %= modulo;

    while exp > 0 {
        if exp % 2 == 1 {
            res = mul_mod(res, base, modulo);
        }
        base = mul_mod(base, base, modulo);
        exp >>= 1;
    }

    res
}

#[inline]
fn mul_mod(mut x: u128, mut y: u128, modu: u128) -> u128 {
    if modu < x { x %= modu; }
    if modu < y { y %= modu; }
    
    let mut res = 0;

    while y > 0 {
        if y & 1 == 1 {
            res = add_mod(res, x, modu);
        }

        y >>= 1;
        x = add_mod(x, x, modu);
    }

    res
}

#[inline]
fn add_mod(mut x: u128, mut y: u128, modu: u128) -> u128 {
    if modu < x { x %= modu; }
    if modu < y { y %= modu; }
    
    if x < modu - y {
        x + y
    } else {
        x.min(y) - (modu - x.max(y))
    }
}

#[inline]
fn gcd(mut m: u128, mut n: u128) -> u128 {
    if m == 0 || n == 0 {
        return m | n;
    }

    let shift = (m | n).trailing_zeros();

    m >>= m.trailing_zeros();
    n >>= n.trailing_zeros();

    while m != n {
        if m > n {
            m -= n;
            m >>= m.trailing_zeros();
        } else {
            n -= m;
            n >>= n.trailing_zeros();
        }
    }
    m << shift
}